import pandas as pd
import tensorflow as tf
import numpy as np

def main():
    # 训练数据
    train_data = pd.read_csv("./csv/run/train_data.csv")
    # print(train_data.head())
    train_data = train_data.drop(train_data.columns[[0]], axis=1)

    # mean = train_data.mean(axis=0)
    # std = train_data.std(axis=0)

    # train_data = (train_data - mean) / std
    
    # 添加时间步长维度
    train_data = train_data.to_numpy()
    train_data = np.repeat(train_data[:, np.newaxis, :], 1, axis=1)
   

    # 标签
    train_label = pd.read_csv("./csv/run/train_label.csv")
    # print(train_label.head())
    train_label = train_label.drop(train_label.columns[[0]], axis=1)

    # 验证数据
    verify_data = pd.read_csv("./csv/run/verify_data.csv")
    # print(verify_data.head())
    verify_data = verify_data.drop(verify_data.columns[[0]], axis=1)
    # verify_data = (verify_data - mean) / std
    # 添加时间步长维度
    verify_data = verify_data.to_numpy()
    verify_data = np.repeat(verify_data[:, np.newaxis, :], 1, axis=1)

    # 标签
    verify_label = pd.read_csv("./csv/run/verify_label.csv")
    # print(verify_label.head())
    verify_label = verify_label.drop(verify_label.columns[[0]], axis=1)


    model = tf.keras.models.Sequential([
        tf.keras.layers.LSTM(64),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(4),
    ])

    # 编译模型
    model.compile(optimizer='rmsprop', loss='mean_squared_error', metrics=['mae'])

    # 训练模型
    model.fit(train_data, train_label, epochs=100)

    # 评估模型
    model.evaluate(verify_data, verify_label)
    # # 测试数据
    # test_data = pd.read_csv("./csv/run/test_data.csv")
    # # print(train_data.head())
    # test_data = test_data.drop(test_data.columns[[0]], axis=1)
    # # test_data = (test_data - mean) / std
    # test_data = test_data.to_numpy()
    # test_data = np.repeat(test_data[:, np.newaxis, :], 1, axis=1)
    # predictions =  model.predict(test_data)
    # np.savetxt("./csv/data_lstm.csv", predictions, delimiter=",")


if __name__ == '__main__':
    main()
